package com.qf.bigdata.repository;

import com.alibaba.fastjson.JSONObject;
import org.dmg.pmml.FieldName;
import org.jpmml.evaluator.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.ClassPathResource;
import org.springframework.stereotype.Component;
import org.xml.sax.SAXException;

import javax.annotation.PostConstruct;
import javax.xml.bind.JAXBException;
import java.io.IOException;
import java.io.InputStream;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

@Component
public class LRModelPredict {
    private static Logger logger= LoggerFactory.getLogger(LRModelPredict.class);
    private static Evaluator evaluator; //解析pmml文件使用


    @PostConstruct //放入IoC容器后立即执行
    public  void init(){
        ClassPathResource classPathResource = new ClassPathResource("lr.pmml");
        try {
            InputStream inputStream = classPathResource.getInputStream();
            evaluator=new LoadingModelEvaluatorBuilder().load(inputStream).build();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (JAXBException e) {
            e.printStackTrace();
        } catch (SAXException e) {
            e.printStackTrace();
        }
        evaluator.verify();
    }

    //输入数据"0.1,-0.1,...."
    public Double predictProbability(String features) {
        List<InputField> inputFields = evaluator.getInputFields();
        String[] featuresArr = features.split(",");
        if(inputFields.size()!=featuresArr.length) {
            logger.error(String.format("模型特征输入长度不够，需要%s，现有%s", inputFields.size(), featuresArr.length));
//            return -1d;
        }
        Map<FieldName, FieldValue> arguments = new LinkedHashMap<>();
        int i=0;
        for (InputField inputField:inputFields) {
            FieldName inputFieldName = inputField.getName();
            String name=inputFieldName.getValue();
            Object rawValue = featuresArr[i];
            i++;
            FieldValue inputFieldValue = inputField.prepare(rawValue);
            arguments.put(inputFieldName,inputFieldValue);
        }

        Map<FieldName, ?> results = evaluator.evaluate(arguments);

        //{label=1, pmml(prediction)=1, prediction=1.0, probability(0)=0.4969870806448454, probability(1)=0.5030129193551546}
        Map<String, ?> resultRecord = EvaluatorUtil.decodeAll(results);//label,probability
        logger.info(resultRecord.toString());
        String label = resultRecord.get("label").toString();
        String probability = resultRecord.get(String.format("probability(%s)", label)).toString();
        return Double.valueOf(probability);
    }


    //完成对输入特征的预测
    //输入数据:{"f1":0.1,"f2":-0.1,...}  uf+item_vec+item_embedding
    //返回Double 预测的概率
    public Double predictProbabilityByJson(JSONObject features){
        List<InputField> inputFields = evaluator.getInputFields();
        Map<FieldName, FieldValue> arguments = new LinkedHashMap<>();
        for (InputField inputField:inputFields) {
            FieldName inputFieldName = inputField.getName();
            String name=inputFieldName.getValue();
            double doubleValue = features.getDoubleValue(name);
            FieldValue inputFieldValue = inputField.prepare(doubleValue);
            arguments.put(inputFieldName,inputFieldValue);
        }

        Map<FieldName, ?> results = evaluator.evaluate(arguments);

        Map<String, ?> resultRecord = EvaluatorUtil.decodeAll(results);//label,probability
        logger.info(resultRecord.toString());
        String label = resultRecord.get("label").toString();
        String probability = resultRecord.get(String.format("probability(%s)", label)).toString();
        return Double.valueOf(probability);
    }

    public static void main(String[] args) {
        LRModelPredict lrModelPredict = new LRModelPredict();
        lrModelPredict.init();
        String json="{'f1':0.1,'f2':-0.1}";
        System.out.println(lrModelPredict.predictProbabilityByJson(JSONObject.parseObject(json)));

        String features="0.1,-0.1";
        System.out.println(lrModelPredict.predictProbability(features));
    }
}
